package net.cyue.ort.llm.gui.config;

import net.cyue.ort.llm.gui.FileFilterUtil;

import javax.swing.*;
import java.awt.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;

/**
 * 右侧配置面板，负责文件选择、状态展示与触发模型加载。
 */
public class ModelConfigPanel extends JPanel implements ModelLoader.Listener {

    private static final String DEFAULT_DIR = "E:\\Softwares\\RWKV Runner\\models";

    private final JFrame owner;
    private final ModelLoader modelLoader;
    private final SelectedModelFiles selectedFiles = new SelectedModelFiles();

    private final JComboBox<ModelType> modelTypeComboBox = new JComboBox<>(ModelType.values());
    private final JLabel tokenizerLabel = createFileLabel("tokenizer文件：未选择");
    private final JLabel tokenizerConfigLabel = createFileLabel("tokenizer_config文件：未选择");
    private final JLabel configLabel = createFileLabel("config文件：未选择");
    private final JLabel generationConfigLabel = createFileLabel("generation_config文件：未选择");
    private final JLabel modelLabel = createFileLabel("模型文件：未选择");
    private final JLabel statusLabel = new JLabel("状态：未加载模型");

    private final JButton selectConfigDirButton = new JButton("选择配置文件目录");
    private final JButton selectModelButton = new JButton("选择模型文件");
    private final JButton loadButton = new JButton("加载模型");

    public ModelConfigPanel(JFrame owner, ModelLoader modelLoader) {
        super();
        this.owner = owner;
        this.modelLoader = modelLoader;

        setLayout(new BoxLayout(this, BoxLayout.Y_AXIS));
        setPreferredSize(new Dimension(owner.getWidth() / 3, owner.getHeight()));
        setBorder(BorderFactory.createEmptyBorder(0, 10, 0, 0));
        buildUI();
        bindActions();
    }

    private void buildUI() {
        JPanel modelTypePanel = new JPanel(new FlowLayout(FlowLayout.LEFT));
        modelTypePanel.add(new JLabel("模型类型:"));
        modelTypePanel.add(modelTypeComboBox);

        add(modelTypePanel);
        add(Box.createVerticalStrut(10));
        add(tokenizerLabel);
        add(tokenizerConfigLabel);
        add(configLabel);
        add(generationConfigLabel);
        add(modelLabel);
        add(Box.createVerticalStrut(10));
        add(selectConfigDirButton);
        add(Box.createVerticalStrut(5));
        add(selectModelButton);
        add(Box.createVerticalStrut(10));
        add(loadButton);
        add(Box.createVerticalStrut(10));
        add(statusLabel);
    }

    private void bindActions() {
        selectConfigDirButton.addActionListener(e -> selectConfigDirectory());
        selectModelButton.addActionListener(e -> selectModelFile());
        loadButton.addActionListener(e -> triggerLoad());
    }

    private void selectConfigDirectory() {
        JFileChooser chooser = new JFileChooser(DEFAULT_DIR);
        chooser.setFileSelectionMode(JFileChooser.DIRECTORIES_ONLY);
        if (chooser.showOpenDialog(owner) != JFileChooser.APPROVE_OPTION) {
            return;
        }

        selectedFiles.clearConfigFiles();
        resetConfigLabels();

        Path directory = chooser.getSelectedFile().toPath();
        try (var paths = Files.list(directory)) {
            paths.forEach(path -> assignConfigFile(path, getSelectedModelType()));
        } catch (Exception ex) {
            JOptionPane.showMessageDialog(owner, "扫描目录失败：" + ex.getMessage(),
                    "错误", JOptionPane.ERROR_MESSAGE);
        }
        refreshLabels();
    }

    private void selectModelFile() {
        JFileChooser chooser = new JFileChooser(DEFAULT_DIR);
        chooser.setFileFilter(FileFilterUtil.createFileFilter(".onnx"));
        chooser.setFileSelectionMode(JFileChooser.FILES_ONLY);
        if (chooser.showOpenDialog(owner) != JFileChooser.APPROVE_OPTION) {
            return;
        }
        selectedFiles.setModel(chooser.getSelectedFile().toPath());
        refreshLabels();
    }

    private void assignConfigFile(Path path, ModelType modelType) {
        String fileName = path.getFileName().toString();
        if (modelType == ModelType.TRANSFORMER && Objects.equals(fileName, "tokenizer.json")) {
            selectedFiles.setTokenizer(path);
        } else if (modelType == ModelType.RWKV && Objects.equals(fileName, "vocab.txt")) {
            selectedFiles.setTokenizer(path);
        } else if (Objects.equals(fileName, "tokenizer_config.json")) {
            selectedFiles.setTokenizerConfig(path);
        } else if (Objects.equals(fileName, "config.json")) {
            selectedFiles.setConfig(path);
        } else if (Objects.equals(fileName, "generation_config.json")) {
            selectedFiles.setGenerationConfig(path);
        }
    }

    private void triggerLoad() {
        if (!selectedFiles.isComplete()) {
            onStatusChanged("状态：请选择所有必需文件");
            return;
        }

        setButtonsEnabled(false);
        CompletableFuture<Void> future = modelLoader.loadAsync(
            selectedFiles, 
            getSelectedModelType(),
            this
        );
        future.whenComplete((unused, throwable) ->
                SwingUtilities.invokeLater(() -> setButtonsEnabled(true)));
    }

    private ModelType getSelectedModelType() {
        return (ModelType) modelTypeComboBox.getSelectedItem();
    }

    private void refreshLabels() {
        tokenizerLabel.setText(selectedFiles.getTokenizer() == null
                ? "tokenizer文件：未选择" : "tokenizer文件：已选择");
        tokenizerConfigLabel.setText(selectedFiles.getTokenizerConfig() == null
                ? "tokenizer_config文件：未选择" : "tokenizer_config文件：已选择");
        configLabel.setText(selectedFiles.getConfig() == null
                ? "config文件：未选择" : "config文件：已选择");
        generationConfigLabel.setText(selectedFiles.getGenerationConfig() == null
                ? "generation_config文件：未选择" : "generation_config文件：已选择");
        modelLabel.setText(selectedFiles.getModel() == null
                ? "模型文件：未选择" : "模型文件：已选择");
    }

    private void resetConfigLabels() {
        tokenizerLabel.setText("tokenizer文件：未选择");
        tokenizerConfigLabel.setText("tokenizer_config文件：未选择");
        configLabel.setText("config文件：未选择");
        generationConfigLabel.setText("generation_config文件：未选择");
    }

    private void setButtonsEnabled(boolean enabled) {
        selectConfigDirButton.setEnabled(enabled);
        selectModelButton.setEnabled(enabled);
        loadButton.setEnabled(enabled);
    }

    private JLabel createFileLabel(String text) {
        JLabel label = new JLabel(text);
        label.setAlignmentX(Component.LEFT_ALIGNMENT);
        return label;
    }

    @Override
    public void onStatusChanged(String status) {
        SwingUtilities.invokeLater(() -> statusLabel.setText(status));
    }

    @Override
    public void onFailure(Exception ex) {
        ModelLoader.Listener.super.onFailure(ex);
        SwingUtilities.invokeLater(() ->
                JOptionPane.showMessageDialog(owner, ex.getMessage(), "模型加载失败",
                        JOptionPane.ERROR_MESSAGE));
    }
}

